import requests
import threading
import os
import time
from util import Logger
from concurrent.futures import wait
from requests.packages import urllib3

# 忽略HTTPS告警
urllib3.disable_warnings()
g_mutex = threading.Lock()
# ⽇志⽂件⽬录为当前⽬录所在⽂件下的log⽬录
log_dir = os.path.dirname(os.path.abspath(__file__))+'/log/'
# ⽬录不存在则创建
if not os.path.exists(log_dir):
    os.mkdir(log_dir)
cur_time = time.strftime('%Y-%m-%d', time.localtime(time.time())).replace(':','').replace(' ', '_')
LogFile = log_dir + cur_time + '.log'
logging = Logger(LogFile, level='debug').logger

class Worker(object):
    site = None
    cfg = None

    def __init__(self, site, cfg, pool):
        self.site = site
        self.cfg = cfg
        self.pool = pool

    def start(self):
            try:
                scan_list = []
                for web_dir in self.cfg['dics']:
                    t = Scanner(self.site, web_dir)
                    future = self.pool.submit(t.run())
                    scan_list.append(future)
                wait(scan_list)
            except Exception as e:
                logging.error('Start error: %s' % e)

class Scanner():
    site = None
    dics = None
    request_method = None
    headers = {
        'user-agent': 'Mozilla/5.0 (Windows NT 6.1; WOW64) '
        'AppleWebKit/537.36 (KHTML, like Gecko) Chrome/43.0.2357.132Safari/537.36'
    }

    def __init__(self, site, dics):
        try:
            self.site = site
            self.dics = dics
            if self.site.endswith('/'):
                self.site = self.site[:-1]
            if '://' not in self.site:
                self.site = 'http://' + self.site
        except Exception as e:
            logging.error('Scanner_init error: %s' % e)
        
    def run(self):
        self.ScanOne(self.site, self.dics)

    def ScanOne(self, site, dic):
        try:
            if not dic.startswith('/'):
                dic = '/' + dic
                url = site + dic
                # 不允许重定向，防⽌重定向误识别为存在的⻚⾯
                res = requests.get(url, verify=False, allow_redirects=False,
                headers=self.headers, timeout=8)
                if 200 == res.status_code:
                    print('\033[93m[+]\033[96m%d\033[0m %s' % (res.status_code, url))
                    self.WriteFile('./out.txt', url)
                else:
                    print('\033[95m[-]\033[31m%d\033[0m %s' % (res.status_code, url))
        except Exception as e:
            if 'Max retries exceeded with url' in str(e):
                pass
            else:
                logging.error('Start error: %s , url:%s' % (e, url))

    def WriteFile(self, file, msg):
        try:
        # 写⽂件时加锁
            g_mutex.acquire()
            with open(file, 'a+') as f:
                f.write('%s\r\n' % (str(msg)))
        except Exception as e:
            logging.error('WriteFile error : %s' % e)